from groundingdino.util.inference import load_model, load_image, predict, annotate
import time
import cv2

model = load_model("groundingdino/config/GroundingDINO_SwinT_OGC.py", "weights/groundingdino_swint_ogc.pth")
IMAGE_PATH = "images/20240611-100727.png"
#TEXT_PROMPT = "tree . road . car . cone ."
TEXT_PROMPT = "car . cone ."
BOX_TRESHOLD = 0.35
TEXT_TRESHOLD = 0.25

image_source, image = load_image(IMAGE_PATH)
start_time = time.time()
boxes, logits, phrases = predict(
    model=model,
    image=image,
    caption=TEXT_PROMPT,
    box_threshold=BOX_TRESHOLD,
    text_threshold=TEXT_TRESHOLD
)
end_time = time.time()
elapsed_time = end_time - start_time
print("elapsed time", elapsed_time)
annotated_frame = annotate(image_source=image_source, boxes=boxes, logits=logits, phrases=phrases)
cv2.imwrite("./annotated_image.jpg", annotated_frame)
